package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"sync"
	"time"
	
	"ebbflow/pkg/cbr"
	"ebbflow/src/consensus"
)

// PerformanceRecord 记录单个操作的性能数据
type PerformanceRecord struct {
	Operation     string        `json:"operation"`
	Epoch         int           `json:"epoch,omitempty"`
	Duration      time.Duration `json:"duration_ms"`
	Success       bool          `json:"success"`
	Timestamp     time.Time     `json:"timestamp"`
}

// CBRPerformanceMetrics 记录 CBR 性能指标
type CBRPerformanceMetrics struct {
	TotalExecutions      int             `json:"total_executions"`
	SuccessfulExecutions int             `json:"successful_executions"`
	FailedExecutions     int             `json:"failed_executions"`
	TotalDuration        time.Duration   `json:"total_duration_ms"`
	AverageDuration      time.Duration   `json:"average_duration_ms"`
	MinDuration          time.Duration    `json:"min_duration_ms"`
	MaxDuration          time.Duration    `json:"max_duration_ms"`
	Executions           []PerformanceRecord `json:"executions"`
}

// ProofChainPerformanceMetrics 记录 ProofChain 性能指标
type ProofChainPerformanceMetrics struct {
	TotalAdditions       int             `json:"total_additions"`
	TotalQueries         int             `json:"total_queries"`
	AddDurations         []time.Duration `json:"add_durations_ms"`
	QueryDurations       []time.Duration `json:"query_durations_ms"`
	AverageAddDuration   time.Duration   `json:"average_add_duration_ms"`
	AverageQueryDuration time.Duration   `json:"average_query_duration_ms"`
	MinAddDuration       time.Duration   `json:"min_add_duration_ms"`
	MaxAddDuration       time.Duration   `json:"max_add_duration_ms"`
	MinQueryDuration     time.Duration   `json:"min_query_duration_ms"`
	MaxQueryDuration     time.Duration   `json:"max_query_duration_ms"`
	ChainLength          int             `json:"chain_length"`
}

// CombinedPerformanceReport 综合性能报告
type CombinedPerformanceReport struct {
	TestConfig            TestConfig                 `json:"test_config"`
	CBRMetrics            CBRPerformanceMetrics      `json:"cbr_metrics"`
	ProofChainMetrics     ProofChainPerformanceMetrics `json:"proofchain_metrics"`
	TotalTestDuration     time.Duration              `json:"total_test_duration_ms"`
	Timestamp             time.Time                  `json:"timestamp"`
}

// TestConfig 测试配置
type TestConfig struct {
	NodeID            int `json:"node_id"`
	TotalNodes        int `json:"total_nodes"`
	FaultyNodes       int `json:"faulty_nodes"`
	Epochs            int `json:"epochs"`
	ReconfigInterval  int `json:"reconfig_interval"`
}

var (
	cbrRecords      []PerformanceRecord
	proofAddRecords []PerformanceRecord
	proofQueryRecords []PerformanceRecord
	recordsMutex    sync.Mutex
)

func main() {
	// Parse command line arguments
	nodeID := flag.Int("id", 0, "Node ID")
	totalNodes := flag.Int("n", 4, "Total number of nodes")
	faultyNodes := flag.Int("f", 1, "Number of faulty nodes")
	epochs := flag.Int("epochs", 20, "Number of epochs to run")
	reconfigInterval := flag.Int("r", 5, "Reconfiguration interval (execute CBR every R epochs)")
	outputFile := flag.String("output", "cbr_performance.json", "Output file for performance data")
	
	flag.Parse()
	
	// Setup logger
	logger := log.New(os.Stdout, fmt.Sprintf("[Node %d] ", *nodeID), log.LstdFlags)
	
	logger.Printf("=== CBR 和 ProofChain 性能测试 ===")
	logger.Printf("配置: N=%d, F=%d, 轮次数=%d, R=%d", *totalNodes, *faultyNodes, *epochs, *reconfigInterval)
	logger.Printf("")
	
	// Initialize performance tracking
	cbrRecords = make([]PerformanceRecord, 0)
	proofAddRecords = make([]PerformanceRecord, 0)
	proofQueryRecords = make([]PerformanceRecord, 0)
	
	// Create CBR instance
	cbrNode := consensus.NewEbbFlowCBR(*nodeID, *totalNodes, *faultyNodes, logger)
	
	testStartTime := time.Now()
	
	// Run CBR for multiple epochs
	logger.Printf("开始运行 CBR 协议和性能测试...")
	logger.Printf("")
	
	for epoch := 0; epoch < *epochs; epoch++ {
		logger.Printf("--- Epoch %d ---", epoch)
		
		// Check if CBR should execute at this epoch
		if cbrNode.ShouldExecute(epoch) {
			logger.Printf("✓ 在 Epoch %d 执行 CBR 重配置", epoch)
			
			// Generate new membership
			newMembership := generateMembership(*totalNodes)
			
			// Measure CBR execution time
			cbrStartTime := time.Now()
			result, err := cbrNode.ExecuteReconfiguration(epoch, newMembership)
			cbrDuration := time.Since(cbrStartTime)
			
			// Record CBR performance
			cbrRecord := PerformanceRecord{
				Operation: "CBR_Execute",
				Epoch:     epoch,
				Duration:  cbrDuration,
				Success:   err == nil,
				Timestamp: time.Now(),
			}
			recordsMutex.Lock()
			cbrRecords = append(cbrRecords, cbrRecord)
			recordsMutex.Unlock()
			
			if err != nil {
				logger.Printf("✗ CBR 执行失败: %v", err)
			} else {
				logger.Printf("✓ CBR 执行成功，耗时: %v", cbrDuration)
				
				// Measure ProofChain query performance (query the entry we just added)
				proofQueryStart := time.Now()
				proofEntry := cbrNode.GetProofEntry(epoch)
				proofQueryDuration := time.Since(proofQueryStart)
				
				proofQueryRecord := PerformanceRecord{
					Operation: "ProofChain_Query",
					Epoch:     epoch,
					Duration:  proofQueryDuration,
					Success:   proofEntry != nil,
					Timestamp: time.Now(),
				}
				recordsMutex.Lock()
				proofQueryRecords = append(proofQueryRecords, proofQueryRecord)
				recordsMutex.Unlock()
				
				if proofEntry != nil {
					logger.Printf("  ProofChain 查询成功，耗时: %v", proofQueryDuration)
					
					// Measure ProofChain addition performance by testing AddEntry separately
					// Create a test entry and measure addition time
					testProofChain := cbr.NewProofChain()
					testEntry := &cbr.ProofEntry{
						Epoch:   epoch + 10000, // Use offset to avoid conflicts
						N:       *totalNodes,
						F:       *faultyNodes,
						L:       len(newMembership),
						Members: make([]int, len(newMembership)),
					}
					for i := range newMembership {
						testEntry.Members[i] = i
					}
					
					addStart := time.Now()
					addErr := testProofChain.AddEntry(testEntry)
					addDuration := time.Since(addStart)
					
					addRecord := PerformanceRecord{
						Operation: "ProofChain_Add",
						Epoch:     epoch,
						Duration:  addDuration,
						Success:   addErr == nil,
						Timestamp: time.Now(),
					}
					recordsMutex.Lock()
					proofAddRecords = append(proofAddRecords, addRecord)
					recordsMutex.Unlock()
					
					if addErr == nil {
						logger.Printf("  ProofChain 添加测试，耗时: %v", addDuration)
					}
				}
				
				logger.Printf("  新成员列表: %v", result.NewMembership)
			}
		} else {
			logger.Printf("  跳过 CBR (非重配置轮次)")
		}
		
		// Periodically query proof chain to measure query performance
		if epoch%3 == 0 && epoch > 0 {
			proofChain := cbrNode.GetProofChain()
			if proofChain != nil && proofChain.Length() > 0 {
				// Query random epochs from proof chain
				for queryEpoch := 0; queryEpoch < epoch; queryEpoch += *reconfigInterval {
					if queryEpoch > 0 {
						queryStart := time.Now()
						entry := cbrNode.GetProofEntry(queryEpoch)
						queryDuration := time.Since(queryStart)
						
						queryRecord := PerformanceRecord{
							Operation: "ProofChain_Query",
							Epoch:     queryEpoch,
							Duration:  queryDuration,
							Success:   entry != nil,
							Timestamp: time.Now(),
						}
						recordsMutex.Lock()
						proofQueryRecords = append(proofQueryRecords, queryRecord)
						recordsMutex.Unlock()
					}
				}
			}
		}
		
		logger.Printf("")
		time.Sleep(50 * time.Millisecond) // 模拟轮次间隔
	}
	
	totalTestDuration := time.Since(testStartTime)
	
	// Get final metrics
	cbrMetrics := cbrNode.GetMetrics()
	proofChain := cbrNode.GetProofChain()
	
	// Calculate performance metrics
	cbrPerfMetrics := calculateCBRPerformanceMetrics(cbrMetrics, cbrRecords)
	proofChainPerfMetrics := calculateProofChainPerformanceMetrics(proofChain, proofAddRecords, proofQueryRecords)
	
	// Create combined report
	report := CombinedPerformanceReport{
		TestConfig: TestConfig{
			NodeID:           *nodeID,
			TotalNodes:       *totalNodes,
			FaultyNodes:      *faultyNodes,
			Epochs:           *epochs,
			ReconfigInterval: *reconfigInterval,
		},
		CBRMetrics:        cbrPerfMetrics,
		ProofChainMetrics: proofChainPerfMetrics,
		TotalTestDuration: totalTestDuration,
		Timestamp:         time.Now(),
	}
	
	// Print summary
	logger.Printf("=== 性能测试总结 ===")
	logger.Printf("")
	logger.Printf("CBR 性能:")
	logger.Printf("  总执行次数: %d", cbrPerfMetrics.TotalExecutions)
	logger.Printf("  成功次数: %d", cbrPerfMetrics.SuccessfulExecutions)
	logger.Printf("  失败次数: %d", cbrPerfMetrics.FailedExecutions)
	if cbrPerfMetrics.TotalExecutions > 0 {
		logger.Printf("  平均耗时: %v", cbrPerfMetrics.AverageDuration)
		logger.Printf("  最小耗时: %v", cbrPerfMetrics.MinDuration)
		logger.Printf("  最大耗时: %v", cbrPerfMetrics.MaxDuration)
	}
	logger.Printf("")
	logger.Printf("ProofChain 性能:")
	logger.Printf("  总添加次数: %d", proofChainPerfMetrics.TotalAdditions)
	logger.Printf("  总查询次数: %d", proofChainPerfMetrics.TotalQueries)
	if proofChainPerfMetrics.TotalAdditions > 0 {
		logger.Printf("  平均添加耗时: %v", proofChainPerfMetrics.AverageAddDuration)
		logger.Printf("  最小添加耗时: %v", proofChainPerfMetrics.MinAddDuration)
		logger.Printf("  最大添加耗时: %v", proofChainPerfMetrics.MaxAddDuration)
	}
	if proofChainPerfMetrics.TotalQueries > 0 {
		logger.Printf("  平均查询耗时: %v", proofChainPerfMetrics.AverageQueryDuration)
		logger.Printf("  最小查询耗时: %v", proofChainPerfMetrics.MinQueryDuration)
		logger.Printf("  最大查询耗时: %v", proofChainPerfMetrics.MaxQueryDuration)
	}
	logger.Printf("  ProofChain 长度: %d", proofChainPerfMetrics.ChainLength)
	logger.Printf("")
	logger.Printf("总测试时间: %v", totalTestDuration)
	logger.Printf("")
	
	// Save to JSON file
	reportJSON, err := json.MarshalIndent(report, "", "  ")
	if err != nil {
		logger.Printf("✗ 序列化性能报告失败: %v", err)
		return
	}
	
	err = os.WriteFile(*outputFile, reportJSON, 0644)
	if err != nil {
		logger.Printf("✗ 写入性能报告失败: %v", err)
		return
	}
	
	logger.Printf("✓ 性能报告已保存到: %s", *outputFile)
	logger.Printf("=== 性能测试完成 ===")
}

func calculateCBRPerformanceMetrics(metrics *cbr.CBRMetrics, records []PerformanceRecord) CBRPerformanceMetrics {
	result := CBRPerformanceMetrics{
		TotalExecutions:      metrics.TotalExecutions,
		SuccessfulExecutions: metrics.SuccessfulExecutions,
		FailedExecutions:     metrics.FailedExecutions,
		TotalDuration:        metrics.TotalDuration,
		AverageDuration:      metrics.AverageDuration,
		Executions:           records,
	}
	
	if len(records) > 0 {
		result.MinDuration = records[0].Duration
		result.MaxDuration = records[0].Duration
		for _, record := range records {
			if record.Duration < result.MinDuration {
				result.MinDuration = record.Duration
			}
			if record.Duration > result.MaxDuration {
				result.MaxDuration = record.Duration
			}
		}
	}
	
	return result
}

func calculateProofChainPerformanceMetrics(chain *cbr.ProofChain, addRecords []PerformanceRecord, queryRecords []PerformanceRecord) ProofChainPerformanceMetrics {
	result := ProofChainPerformanceMetrics{
		TotalAdditions: len(addRecords),
		TotalQueries:   len(queryRecords),
		ChainLength:    0,
	}
	
	if chain != nil {
		result.ChainLength = chain.Length()
	}
	
	// Extract durations from records
	addDurations := make([]time.Duration, 0, len(addRecords))
	for _, record := range addRecords {
		if record.Success {
			addDurations = append(addDurations, record.Duration)
		}
	}
	
	queryDurations := make([]time.Duration, 0, len(queryRecords))
	for _, record := range queryRecords {
		if record.Success {
			queryDurations = append(queryDurations, record.Duration)
		}
	}
	
	result.AddDurations = addDurations
	result.QueryDurations = queryDurations
	
	// Calculate statistics
	if len(addDurations) > 0 {
		var total time.Duration
		result.MinAddDuration = addDurations[0]
		result.MaxAddDuration = addDurations[0]
		for _, d := range addDurations {
			total += d
			if d < result.MinAddDuration {
				result.MinAddDuration = d
			}
			if d > result.MaxAddDuration {
				result.MaxAddDuration = d
			}
		}
		result.AverageAddDuration = total / time.Duration(len(addDurations))
	}
	
	if len(queryDurations) > 0 {
		var total time.Duration
		result.MinQueryDuration = queryDurations[0]
		result.MaxQueryDuration = queryDurations[0]
		for _, d := range queryDurations {
			total += d
			if d < result.MinQueryDuration {
				result.MinQueryDuration = d
			}
			if d > result.MaxQueryDuration {
				result.MaxQueryDuration = d
			}
		}
		result.AverageQueryDuration = total / time.Duration(len(queryDurations))
	}
	
	return result
}

func generateMembership(totalNodes int) []string {
	membership := make([]string, totalNodes)
	for i := 0; i < totalNodes; i++ {
		membership[i] = fmt.Sprintf("127.0.0.1:%d", 8000+i)
	}
	return membership
}

